"""
1.	利用Pytorch深度学习框架，建立由卷积神经网络Resnet34构造的特征提取主干模型，和后端分类网络，进行猫狗图像识别，并进行模型评估，
按下面要求完成相应代码（75分）
"""
import torch as pt
import torchvision as ptv
import numpy as np
import os
import sys
import cv2 as cv
from sklearn.model_selection import train_test_split
from torch.utils.data import TensorDataset, DataLoader
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, precision_score, recall_score, f1_score
import seaborn as sns

np.random.seed(1)
pt.random.manual_seed(1)

# ②	定义训练批次、周期、图像尺寸、学习率等超参数
VER = 'v1.0'
SIZE = 224
BATCH_SIZE = 16
N_EPOCHS = 2
ALPHA = 1e-4
device = 'cuda:0' if pt.cuda.is_available() else 'cpu'
print('device', device)
device = pt.device(device)

# ①	读入猫狗图像数据集，进行适当训练集、验证集、测试集比例划分
print('Reading data ...')
BASE_DIR, FILE_NAME = os.path.split(__file__)
dir = '../../../../large_data/DL1/_many_files/cats_and_dogs_filtered/train'
IMG_DIR = os.path.join(BASE_DIR, dir)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
SAVE_PATH = os.path.join(SAVE_DIR, 'weights.pth')
idx2label, label2idx = {}, {}
yi = 0
x, y = [], []
for sub_dir_name in os.listdir(IMG_DIR):
    sub_dir_path = os.path.join(IMG_DIR, sub_dir_name)
    if not os.path.isdir(sub_dir_path):
        continue
    idx2label[yi] = sub_dir_name
    label2idx[sub_dir_name] = yi
    for file_name in os.listdir(sub_dir_path):
        basename, ext = os.path.splitext(file_name)
        ext = ext.lower()
        if '.jpg' != ext:
            continue
        file_path = os.path.join(sub_dir_path, file_name)
        img = cv.imread(file_path, cv.IMREAD_COLOR)
        img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
        img = cv.resize(img, (SIZE, SIZE), interpolation=cv.INTER_CUBIC)
        x.append(img)
        y.append(yi)
    yi += 1
x = np.uint8(x)
y = np.asarray(y, dtype=np.int64)
x_ori = x.copy()
x = np.float32(x) / 255. * 2. - 1.
x = x.transpose((0, 3, 1, 2))
n_cls = len(np.unique(y))
print('n_cls', n_cls)
print('x', x.shape, x.dtype)
print('y', y.shape, y.dtype)
print('x_ori', x_ori.shape, x_ori.dtype)

x_train, x_test_val, y_train, y_test_val, x_ori_train, x_ori_test_val = train_test_split(
    x, y, x_ori, train_size=0.8, random_state=2, shuffle=True,
)
x_test, x_val, y_test, y_val, x_ori_test, x_ori_val = train_test_split(
    x_test_val, y_test_val, x_ori_test_val, train_size=0.5, random_state=1, shuffle=True,
)
print('x_train', x_train.shape)
print('x_test', x_test.shape)
print('x_val', x_val.shape)
print('y_train', y_train.shape)
print('y_test', y_test.shape)
print('y_val', y_val.shape)
print('x_ori_train', x_ori_train.shape)
print('x_ori_test', x_ori_test.shape)
print('x_ori_val', x_ori_val.shape)

x_train = pt.from_numpy(x_train)
x_test = pt.from_numpy(x_test)
x_val = pt.from_numpy(x_val)
y_train = pt.from_numpy(y_train)
y_test = pt.from_numpy(y_test)
y_val = pt.from_numpy(y_val)

ds_train = TensorDataset(x_train, y_train)
ds_test = TensorDataset(x_test, y_test)
ds_val = TensorDataset(x_val, y_val)

dl_train = DataLoader(ds_train, BATCH_SIZE)
dl_test = DataLoader(ds_test, BATCH_SIZE)
dl_val = DataLoader(ds_val, BATCH_SIZE)

# ③	导入Resnet34主干网络作为预训练模型
resnet34 = ptv.models.resnet34(pretrained=True)
for param in resnet34.parameters():
    param.requires_grad = False
base_model_out_ch = resnet34.fc.in_features
base_model = pt.nn.Sequential(*(list(resnet34.children())[:-1]))

# ④	搭建后端网络模型
customer_model = pt.nn.Linear(base_model_out_ch, n_cls)


# ⑤	打印输出整体模型维度结构
class Resnet34TransLearningModel(pt.nn.Module):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.base = base_model
        self.customer = customer_model

    def forward(self, x):
        x = self.base(x)
        x = pt.squeeze(x, dim=3)
        x = pt.squeeze(x, dim=2)
        x = self.customer(x)
        return x


model = Resnet34TransLearningModel().to(device)
print(model)

# ⑥	使用适当的损失函数loss和优化器optimizer
criterion = pt.nn.CrossEntropyLoss()
optim = pt.optim.Adam(model.parameters(), lr=ALPHA)


# ⑦	定义计算模型准确率函数
def accuracy(y_true, y_pred):
    y_true = y_true.long()
    y_pred = y_pred.argmax(dim=1)
    acc = pt.eq(y_true, y_pred).float().mean()
    return acc


# process data
def process_data(dl, label, is_train, epoch, is_return_data=False):
    avg_loss, avg_acc = 0., 0.
    dl_len = len(dl)
    group = int(np.ceil(dl_len / 10))
    y_true, y_pred = None, None
    for i, (bx, by) in enumerate(dl):
        bx = bx.float().to(device)
        by = by.long().to(device)
        if is_train:
            model.train(True)
            optim.zero_grad()
            pred = model(bx)
            loss = criterion(pred, by)
            loss.backward()
            optim.step()
            model.train(False)
        else:
            model.train(False)
            pred = model(bx)
            loss = criterion(pred, by)

        acc = accuracy(by, pred)
        loss = loss.detach().cpu().numpy()
        acc = acc.detach().cpu().numpy()
        avg_loss += loss
        avg_acc += acc
        if i % group == 0 or i == dl_len - 1:
            print(f'{label}: epoch#{epoch}: batch#{i + 1}: loss = {loss}, acc = {acc}')

        if is_return_data:
            by = by.detach().cpu().numpy().astype(np.int64)
            pred = pred.argmax(dim=1).detach().cpu().numpy().astype(np.int64)
            if y_true is None:
                y_true = by
                y_pred = pred
            else:
                y_true = np.concatenate((y_true, by), axis=0, dtype=np.int64)
                y_pred = np.concatenate((y_pred, pred), axis=0, dtype=np.int64)

    avg_loss /= i + 1
    avg_acc /= i + 1
    return avg_loss, avg_acc, y_true, y_pred


    # ⑧	打印输出每一训练周期训练集和验证集的平均损失值和平均准确率
if os.path.exists(SAVE_DIR):
    print('Loading ...')
    dict = pt.load(SAVE_PATH)
    model.load_state_dict(dict)
    print('Loaded')
    spr = 1
    spc = 2
    spn = 0
    plt.figure(figsize=[12, 6])
else:
    loss_his, acc_his, loss_his_val, acc_his_val = [], [], [], []
    for epoch in range(N_EPOCHS):
        avg_loss, avg_acc, _, _ = process_data(dl_train, 'train', True, epoch + 1)
        print(f'train: epoch#{epoch+1}: avg_loss = {avg_loss}, avg_acc = {avg_acc}')
        avg_loss_val, avg_acc_val, _, _ = process_data(dl_val, 'val', False, epoch + 1)
        print(f'train: epoch#{epoch+1}: avg_loss = {avg_loss}, avg_acc = {avg_acc}, avg_loss_val = {avg_loss_val}, avg_acc_val = {avg_acc_val}')
        loss_his.append(avg_loss)
        acc_his.append(avg_acc)
        loss_his_val.append(avg_loss_val)
        acc_his_val.append(avg_acc_val)
    print('Saving ...')
    os.makedirs(SAVE_DIR, exist_ok=True)
    dict = model.state_dict()
    pt.save(dict, SAVE_PATH)
    print('Saved.')

    # ⑨	图示输出训练集和验证集的平均损失值曲线
    spr = 2
    spc = 2
    spn = 0
    plt.figure(figsize=[6, 6])

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('loss')
    plt.plot(loss_his, label='train')
    plt.plot(loss_his_val, label='val')
    plt.grid()
    plt.legend()

    # ⑩	图示输出训练集和验证集的平均准确率曲线
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('acc')
    plt.plot(acc_his, label='train')
    plt.plot(acc_his_val, label='val')
    plt.grid()
    plt.legend()

# 11	模型训练完成后，用测试集测试模型猫狗分类正确与不正确百分比的饼图
avg_loss, avg_acc, y_true, y_pred = process_data(dl_test, 'test', False, 1, is_return_data=True)
print(f'Test: avg_loss = {avg_loss}, avg_acc = {avg_acc}')
spn += 1
plt.subplot(spr, spc, spn)
plt.title('accuracy on test')
plt.pie([avg_acc, 1-avg_acc], explode=[0.1, 0], labels=['acc', ''], autopct='%.2f%%')

# 12	用测试集数据，打印输出混淆矩阵，并用热力图显示
spn += 1
plt.subplot(spr, spc, spn)
mat = confusion_matrix(y_true, y_pred)
sns.heatmap(mat, annot=True)
print('Please check and close the plotting window to continue ...')
plt.show()

# 13	打印输出查准率P和查全率R
p = precision_score(y_true, y_pred)
r = recall_score(y_true, y_pred)
print(f'查准率P={p}和查全率R={r}')

# 14	打印输出准确率acc和F1分值
f1 = f1_score(y_true, y_pred)
print(f'准确率acc={avg_acc}和F1分值={f1}')

# 15	从测试集中随机选取9张图像，图示输出预测种类名称和真实种类名称，进行比较（参考下图）
spr = 3
spc = 3
spn = 0
plt.figure(figsize=[6, 6])
for i in range(spr * spc):  # ATTENTION: spc wrongly be written as spn, so no image is shown.
    spn += 1
    plt.subplot(spr, spc, spn)
    y = y_true[i]
    pred = y_pred[i]
    judge = 'V' if y == pred else 'X'
    title = f'{idx2label[y]}=>{idx2label[pred]} ({judge})'
    plt.title(title, color='black' if y == pred else 'red')
    plt.axis('off')
    img = x_ori_test[i]
    plt.imshow(img)
print('Please check and close the plotting window to finish ...')
plt.show()
print('Over')
